#!/usr/bin/python2.7
# Copyright 2020 Makani Technologies LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Automatically generate C code for loading JSON data into C structures.

Throughout a "type" refers to a ctypes data structure generated by h2py, i.e:
  - fundamental types (ctypes.c_double, ctypes.c_float, ...),
  - strings (arrays of ctypes.c_char),
  - arrays,
  - structures.
For arrays and structures, the function GetLoadFunc generates the
prototype and definition of the corresponding JSON loading functions.
Manual implementations of the JSON loading functions for fundamental
and string types should be provided in json_load_basic.{c,h}.

Given a type, GetTypeName returns the "type name" used to uniquely
identify the loading function.  Note that ctypes treats arrays with
the same number of indices but different lengths as distinct types,
but we assign these all the same loading function and type name.

GetRequiredTypes can be used to get a recursively generated set of types whose
loading functions are required by a type's loading function.

TODO: Add unit tests for the functions in this file.
"""

import ctypes
from makani.lib.python.autogen import autogen_util


def GetTypeName(t):
  """Return the name used to represent a type in function definitions.

  The ctypes library treats arrays of different length as distinct types,
  whereas this library treats them the same (i.e. double foo[4][5] is mapped
  to 'Array2D_Double').

  Args:
    t: A type.

  Returns:
    A string (e.g. 'Array1D_Double' or 'Struct_CalParams') that represents a
    type.
  """
  return _GetTypeInfo(t)[0]


def GetRequiredTypes(t):
  """Get types whose loading functions are required to load a given type.

  Args:
    t: A type.

  Returns:
    A set of types whose loading functions are required to implement
    the loading function for t.
  """
  types = set()
  if _IsArrayType(t):
    scalar_type = _ArrayCTypeScalarType(t)
    types = {scalar_type}
    types = types.union(GetRequiredTypes(scalar_type))
  elif _IsStructType(t):
    fields = autogen_util.GetCFields(t)
    types = {f_type for (_, f_type) in fields}
    types = types.union(*[GetRequiredTypes(f_type) for (_, f_type) in fields])

  return types


def GetLoadFunc(t):
  """Generate the JSON loading implementation for a given type.

  Args:
    t: A type.

  Returns:
    A pair containing C code.  The first element is a string with the prototype
    of the function and the second element is a list of strings giving the lines
    of C code.

  Raises:
    TypeError: The given type is a string type, a fundamental type, or unknown.
  """
  prototype = _GetLoadFunctionPrototype(t)
  body = []  # This will be an array strings corresponding to lines of code.

  if _IsArrayType(t):
    (_, _, args) = _GetTypeInfo(t)
    body = ['if (JSONLoadArrayCheck(obj, len0) < 0) return -1;',
            'int32_t status;']

    sub_dim = ['len%d' % d for d in range(1, len(args))]
    if sub_dim:
      # TODO: Check for overflow?
      body += ['uint32_t stride = %s;' % ' * '.join(sub_dim)]
      dst = '&s[i*stride]'
    else:
      dst = 's[i]'

    scalar_type = _ArrayCTypeScalarType(t)
    body += ['for (uint32_t i = 0U; i < len0; ++i) {',
             '  ' + _GetLoadFuncCall(scalar_type, 'json_array_get(obj, i)',
                                     dst, 'status', sub_dim),
             '  if (status == -1) {',
             '    JSONLoadPrintIndexError(i);',
             '    return -1;',
             '  }',
             '}',
             'return 0;']

  elif _IsStructType(t):
    fields = autogen_util.GetCFields(t)
    body = ['if (s == NULL) return -1;',
            'if (JSONLoadStructCheck(obj, %uU) < 0) return -1;' % len(fields),
            'json_t *value;',
            'int32_t status;',
            '']
    for (f_name, f_type) in fields:
      body += ['value = JSONLoadStructGet(obj, "%s");' % f_name,
               'if (value == NULL) return -1;',
               _GetLoadFuncCall(f_type, 'value', 's->%s' % f_name, 'status'),
               'if (status == -1) {',
               '  JSONLoadPrintKeyError("%s");' % f_name,
               '  return -1;',
               '}',
               '']
    body += ['return 0;']

  # These functions are hand written in the files json_load_basic.{c,h}.
  elif _IsFundamentalType(t) or _IsStringType(t):
    raise TypeError('Cannot generate load function for type %s.' % str(t))

  else:
    raise TypeError('Unknown type: %s.' % str(t))

  return (prototype, ['' if not line else '  ' + line for line in body])


def _GetLoadFuncName(pretty_name):
  """Returns the name of the function for loading a given type name."""
  return 'JSONLoad' + pretty_name


def _GetLoadFunctionPrototype(t):
  """Produce the prototype of the JSON loading function for a type.

  Example string:
    'int32_t JSONLoadStruct_CalParams(json_t *obj, CalParams *s)'

  Args:
    t: A type.

  Returns:
    A string containing the C code prototype for the loading function for t.
  """
  (pretty_name, ptr_name, dim) = _GetTypeInfo(t)

  if dim:
    args = ['uint32_t len%d' % i for i in range(0, len(dim))]
  else:
    args = []

  args = ['const json_t *obj'] + args + ['%s *s' % ptr_name]

  func_name = _GetLoadFuncName(pretty_name)

  return 'int32_t %s(%s)' % (func_name, ', '.join(args))


def _GetLoadFuncCall(t, src, dst, status, args=None):
  """Generate the string for a statement calling a JSON load function.

  Args:
    t: Type of object being loaded.
    src: String containing a C expression evaluating to the json_t pointer to be
        loaded.
    dst: String containing a C expression evaluating to the name of destination
        variable for fundamental types and structures, or to the name of a
        destination array or pointer for strings and arrays.
    status: String containing the name of a int32_t C variable that will
        record the success or failure of the loading operation.
    args: List of strings containing C expressions for dimensions when the
        caller is an array and the type being loaded is an array or string.
        Default value is None. When None, arrays and strings use their
        dimensions in the function call.

  Returns:
    A string containing a C statement (including semi-colon) that
    attempts to load the data and assigns status.
  """
  assert (_IsStringType(t) or _IsArrayType(t)) or (args is None or not args)

  (pretty_name, _, dim) = _GetTypeInfo(t)
  func_name = _GetLoadFuncName(pretty_name)

  if _IsFundamentalType(t):
    # Assignment is used for fundamental types to avoid taking references
    # to enum types.  See json_load_basic.{c,h} for implementations.
    return '%s = %s(%s, &%s);' % (dst, func_name, src, status)
  else:
    if not dim:
      # For types without dimensions, the argument is not a pointer so we take
      # a reference.
      dst = '&' + dst
    elif args is None:
      # When args is None and dim is not empty, we have been handed an array
      # that is a  field from a structure.  We take a reference to the first
      # element to avoid compiler warnings about implicit casting from array
      # to pointer.
      dst = '&' + dst + '[0]'*len(dim)

    if args is None:
      # When args is None we provide the length arguments based on the
      # dimensions.
      args = ['%dU' % d for d in dim]
    else:
      assert len(args) == len(dim)

    if args:
      return '%s = %s(%s, %s, %s);' % (status, func_name, src, ', '.join(args),
                                       dst)
    else:
      return '%s = %s(%s, %s);' % (status, func_name, src, dst)


def _ArrayCTypeLength(t):
  return t._length_  # pylint: disable=protected-access


def _ArrayCTypeScalarType(t):
  return t._type_  # pylint: disable=protected-access


def _IsStringType(t):
  """Tests if a type is an array of characters."""
  return (type(t) is type(ctypes.Array)
          and _ArrayCTypeScalarType(t) == ctypes.c_char)


def _IsArrayType(t):
  """Tests if a type is an array and also not a string."""
  return (type(t) is type(ctypes.Array)
          and _ArrayCTypeScalarType(t) != ctypes.c_char)


def _IsStructType(t):
  """Tests if a type is a structure."""
  return type(t) is type(ctypes.Structure)


def _IsFundamentalType(t):
  """Tests if a type is fundamental C type."""
  return type(t) is type(ctypes.c_double)


def _GetArrayTypeInfo(t):
  """Return the element type and dimension of a multi-dimensional array.

  Args:
    t: A type representing an array of non-character elements.

  Returns:
    A pair, the first element being the scalar type stored by this
    array and the second element being a list of array dimensions
    (e.g. double foo[2][2] would be associated with (ctypes.c_double,
    [2 2])).
  """
  assert _IsArrayType(t)

  scalar_type = _ArrayCTypeScalarType(t)
  dim = [_ArrayCTypeLength(t)]

  if _IsArrayType(scalar_type):
    (elt_scalar_type, elt_dim) = _GetArrayTypeInfo(scalar_type)
    return (elt_scalar_type, dim + elt_dim)
  else:
    return (scalar_type, dim)


def _GetTypeInfo(t):
  """Return information about a type.

  Args:
    t: A type.

  Returns:
    A tuple of three elements.  The first element is the type name,
    the second entry is a string for the type of pointer that
    represents this object (e.g. "double" for both Array1D_double and
    Array2D_double and "char" for String), and the third element is a
    list of dimensions indicating the length of strings and arrays.

  Raises:
    TypeError: A type was passed in which this function cannot handle.
  """
  ctype_to_str_dict = {
      ctypes.c_bool: ('Boolean', 'bool'),
      ctypes.c_double: ('Double', 'double'),
      ctypes.c_float: ('Float', 'float'),
      ctypes.c_int8: ('Int8', 'int8_t'),
      ctypes.c_int16: ('Int16', 'int16_t'),
      ctypes.c_int32: ('Int32', 'int32_t'),
      ctypes.c_int64: ('Int64', 'int64_t'),
      ctypes.c_uint8: ('UInt8', 'uint8_t'),
      ctypes.c_uint16: ('UInt16', 'uint16_t'),
      ctypes.c_uint32: ('UInt32', 'uint32_t'),
      ctypes.c_uint64: ('UInt64', 'uint64_t'),
  }

  if _IsStringType(t):
    return ('String', 'char', [_ArrayCTypeLength(t)])
  elif _IsArrayType(t):
    (scalar_type, dim) = _GetArrayTypeInfo(t)
    (scalar_type_name, scalar_ptr_name, scalar_dim) = _GetTypeInfo(scalar_type)
    return ('Array%dD_%s' % (len(dim), scalar_type_name), scalar_ptr_name,
            dim + scalar_dim)
  elif _IsStructType(t):
    return ('Struct_' + autogen_util.CStructName(t),
            autogen_util.CStructName(t), [])
  elif _IsFundamentalType(t) and t in ctype_to_str_dict:
    names = ctype_to_str_dict[t]
    return (names[0], names[1], [])
  else:
    raise TypeError('Unknown type: %s.' % str(t))
